/*
Copyright (C) 2010- Peter Bui and Aaron Dingler.
This software is distributed under the GNU General Public License.
See the file COPYING for details.
*/

#include "smq_common.h"
#include "smq_exchange.h"
#include "smq_message_queue.h"

#include <catalog_server.h>
#include <change_process_title.h>
#include <copy_stream.h>
#include <datagram.h>
#include <debug.h>
#include <domain_name_cache.h>
#include <hash_table.h>
#include <link.h>
#include <list.h>
#include <sort_dir.h>
#include <stringtools.h>
#include <username.h>
#include <xmalloc.h>

#include <errno.h>
#include <libgen.h>
#include <signal.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <sys/stat.h>
#include <sys/statfs.h>
#include <sys/utsname.h>
#include <sys/wait.h>

/* Manager configuration ---------------------------------------------------- */

static char *SMQ_Manager_Root		    = "/tmp/spice-messaging-queue";
static char *SMQ_Manager_Relay_Queue	    = SMQ_DEFAULT_RELAY_QUEUE;
static int   SMQ_Manager_Port		    = SMQ_MANAGER_DEFAULT_PORT;
static int   SMQ_Manager_Max_Connections    = 64;
static int   SMQ_Manager_Total_Connections  = 0;
static int   SMQ_Manager_Update_Alarm	    = 0;
static int   SMQ_Manager_Idle_Timeout	    = 60;
static int   SMQ_Manager_Stalled_Timeout    = 3600;
static int   SMQ_Manager_Update_Timeout	    = 60;
static int   SMQ_Manager_Exchange_Timeout   = 60;

static struct hash_table * SMQ_Manager_Exchange_Table = NULL;

static char  SMQ_Manager_Host[DOMAIN_NAME_MAX]  = "";
static char  SMQ_Manager_Owner[USERNAME_MAX]	= "";
static char  SMQ_Manager_CPU[SMQ_LINE_MAX]	= "";
static char  SMQ_Manager_OpSys[SMQ_LINE_MAX]	= "";
static char  SMQ_Manager_OpSysVersion[SMQ_LINE_MAX]	= "";
static int   SMQ_Manager_Start_Time = 0;
static char  SMQ_Manager_Bin_Path[SMQ_LINE_MAX];

static struct datagram *SMQ_Manager_Catalog_Port = NULL;

/* Command line parsing ----------------------------------------------------- */

static void
smq_manager_show_help(const char *progname)
{
    printf("usage: %s [options]\n", progname);
    printf("Most Common Options:\n");
    printf(" -r <dir>    Root of the queue directory (default is %s).\n", SMQ_Manager_Root);
    printf(" -d <flag>   Enable debugging for this subsystem.\n");
    printf(" -p <port>   Listen on this port. (default is %d)\n", SMQ_Manager_Port);
    printf(" -h          This message.\n");
    printf("\nLess Common Options:\n");
    printf(" -t <time>   Disconnect idle clients after this time. (default is %ds)\n", SMQ_Manager_Idle_Timeout);
    printf(" -U <time>   Send status updates at this interval. (default is %ds)\n", SMQ_Manager_Update_Timeout);
}

static void
smq_manager_parse_command_line_arguments(int argc, char *argv[])
{
    int	c;

    /* TODO: expose more configuration options */

    while ((c = getopt(argc, argv, "d:p:r:t:U:h")) != -1) {
	switch (c) {
	    case 'd': 
		debug_flags_set(optarg); 
		break;
	    case 'p': 
		SMQ_Manager_Port = atoi(optarg);
		break;
	    case 'r': 
		SMQ_Manager_Root = optarg;
		break;
	    case 't':
		SMQ_Manager_Idle_Timeout = string_time_parse(optarg);
		break;
	    case 'U': 
		SMQ_Manager_Update_Timeout = string_time_parse(optarg);
		break;
	    case 'h':
		smq_manager_show_help(argv[0]); 
		exit(EXIT_SUCCESS);
		break;
	    default:
		smq_manager_show_help(argv[0]); 
		exit(EXIT_FAILURE);
		break;
	}
    }
}

/* Setup relay -------------------------------------------------------------- */

static void
smq_manager_setup_relay()
{
    char src_path[SMQ_LINE_MAX];
    char tgt_path[SMQ_LINE_MAX];
    FILE *src_stream = NULL;
    FILE *tgt_stream = NULL;

    if (!smq_message_queue_create(SMQ_Manager_Root, SMQ_Manager_Relay_Queue)) 
	fatal("unable to create relay queue %s/%s: %s", SMQ_Manager_Root, SMQ_Manager_Relay_Queue, strerror(errno));

    snprintf(src_path, SMQ_LINE_MAX, "%s/smq_relay.py", SMQ_Manager_Bin_Path);
    snprintf(tgt_path, SMQ_LINE_MAX, "%s/%s/.bindings/smq_relay.py", SMQ_Manager_Root, SMQ_Manager_Relay_Queue);

    src_stream = fopen(src_path, "r");
    if (!src_stream) fatal("could not setup relay: %s", strerror(errno));

    tgt_stream = fopen(tgt_path, "w+");
    if (!tgt_stream) fatal("could not setup relay: %s", strerror(errno));

    if (copy_stream_to_stream(src_stream, tgt_stream) < 0)
	fatal("could not setup relay: %s", strerror(errno));

    fclose(src_stream);
    fclose(tgt_stream);

    chmod(tgt_path, 0700);
}

/* Start exchange ----------------------------------------------------------- */

static void
smq_manager_start_exchanges(int sig)
{
    char **mq_list;
    pid_t pid;

    if (sort_dir(SMQ_Manager_Root, &mq_list, strcmp)) {
	int i;

	for (i = 0; mq_list[i]; i++) {
            if (mq_list[i][0] != '.') {
		pid = (pid_t)(PTRINT_T)hash_table_lookup(SMQ_Manager_Exchange_Table, mq_list[i]);
		if (pid <= 0) {
		    debug(D_DEBUG, "start_exchange: %s", mq_list[i]);
		    pid = fork();
		    if (pid == 0) {
			setenv("SMQ_MANAGER_HOST", SMQ_Manager_Host, 1);
			setenv("SMQ_MANAGER_BIN_PATH", SMQ_Manager_Bin_Path, 1);
			change_process_title("smq_exchange: %s", mq_list[i]);
			smq_exchange_run(SMQ_Manager_Root, mq_list[i], SMQ_Manager_Exchange_Timeout);
			_exit(0);
		    } else if (pid >= 0) {
			if (!hash_table_insert(SMQ_Manager_Exchange_Table, mq_list[i], (void *)(PTRINT_T)pid))
			    fatal("unable to insert pid to exchange table");
		    } else {
			fatal("unable to start exchange %s: %s", mq_list[i], strerror(errno));
		    }
		}
	    }
        }

	sort_dir_free(mq_list);
    }
}

/* Signal handlers ---------------------------------------------------------- */

static void 
smq_manager_ignore_signal(int sig)
{
}

static void 
smq_manager_shutdown_clean(int sig)
{
    exit(0);
}

static void 
smq_manager_reap_child(int sig)
{
    pid_t pid;
    int status;
    int is_connection = 1;

    do {
	pid = waitpid(-1, &status, WNOHANG);
	if (pid > 0) {
	    char *key;
	    void *value;
	    hash_table_firstkey(SMQ_Manager_Exchange_Table);
	    while (hash_table_nextkey(SMQ_Manager_Exchange_Table, &key, &value)) {
		if (pid == (pid_t)(PTRINT_T)value) {
		    debug(D_DEBUG, "reap %s exchange process %d", key, pid);
		    is_connection = 0;
		    break;
		}
	    }

	    if (is_connection) {
		debug(D_DEBUG, "reap connection process %d", pid);
		SMQ_Manager_Total_Connections--;
	    } else {
		hash_table_remove(SMQ_Manager_Exchange_Table, key);
	    }
	}
    } while (pid > 0);
}

static void
smq_manager_install_signal_handler(int sig, void (*handler)(int sig))
{
    struct sigaction s;

    s.sa_handler = handler;
    sigfillset(&s.sa_mask);
    s.sa_flags = 0;
    sigaction(sig, &s, 0);
}

/* Process client ----------------------------------------------------------- */

static void
smq_manager_process_client(struct link *client_link)
{
    char line[SMQ_LINE_MAX];
    INT64_T line_size;

    link_tune(client_link, LINK_TUNE_INTERACTIVE);

    while (1) {
	INT64_T result;
	int argc;
	char **argv;

	if (!link_readline(client_link, line, SMQ_LINE_MAX, time(0) + SMQ_Manager_Idle_Timeout)) {
	    debug(D_DEBUG, "timeout: client idle too long");
	    break;
	}

	if (!string_split_quotes(line, &argc, &argv)) {
	    debug(D_DEBUG, "unable to split line: %s", line);
	    break;
	}

	if (argc <= 0) {
	    free(argv);
	    continue;
	}

	if (strcmp(argv[0], "create") == 0 && argc == 2) {
	    debug(D_DEBUG, "Create RPC");
	    result = smq_message_queue_create(SMQ_Manager_Root, argv[1]);
	    if (result) kill(getppid(), SIGUSR1);
	    debug(D_DEBUG, "= %lld", result);
	} else if (strcmp(argv[0], "put") == 0 && argc == 4) {
	    debug(D_DEBUG, "Put RPC");
	    int meta_size = strtol(argv[2], NULL, 10);
	    int body_size = strtol(argv[3], NULL, 10);
	    result = smq_message_queue_put(SMQ_Manager_Root, argv[1], meta_size, body_size, client_link, SMQ_Manager_Stalled_Timeout);
	    if (result) {
		pid_t pid = (pid_t)(PTRINT_T)hash_table_lookup(SMQ_Manager_Exchange_Table, argv[1]);
		if (pid) kill(pid, SIGHUP);
	    }
	    debug(D_DEBUG, "= %lld", result);
	} else if (strcmp(argv[0], "get") == 0 && argc == 3) {
	    debug(D_DEBUG, "Get RPC");
	    result = smq_message_queue_get(SMQ_Manager_Root, argv[1], argv[2], client_link, SMQ_Manager_Stalled_Timeout);
	    debug(D_DEBUG, "= %lld", result);
	} else if (strcmp(argv[0], "list") == 0 && argc == 2) {
	    debug(D_DEBUG, "List RPC");
	    result = smq_message_queue_list(SMQ_Manager_Root, argv[1], client_link, SMQ_Manager_Stalled_Timeout);
	    debug(D_DEBUG, "= %lld", result);
	} else if (strcmp(argv[0], "flush") == 0 && argc == 2) {
	    debug(D_DEBUG, "Flush RPC");
	    result = smq_message_queue_flush(SMQ_Manager_Root, argv[1], client_link, SMQ_Manager_Stalled_Timeout);
	    debug(D_DEBUG, "= %lld", result);
	} else if (strcmp(argv[0], "bind") == 0 && argc == 4) {
	    debug(D_DEBUG, "Bind RPC");
	    char *bind_name = argv[2];
	    int   bind_size = strtol(argv[3], NULL, 10);
	    result = smq_message_queue_bind(SMQ_Manager_Root, argv[1], bind_name, bind_size, client_link, SMQ_Manager_Stalled_Timeout);
	    debug(D_DEBUG, "= %lld", result);
	} else if (strcmp(argv[0], "unbind") == 0 && argc == 3) {
	    debug(D_DEBUG, "Unbind RPC");
	    char *bind_name = argv[2];
	    result = smq_message_queue_unbind(SMQ_Manager_Root, argv[1], bind_name, client_link, SMQ_Manager_Stalled_Timeout);
	    debug(D_DEBUG, "= %lld", result);
	} else if (strcmp(argv[0], "bindings") == 0 && argc == 2) {
	    debug(D_DEBUG, "Bindings RPC");
	    result = smq_message_queue_bindings(SMQ_Manager_Root, argv[1], client_link, SMQ_Manager_Stalled_Timeout);
	    debug(D_DEBUG, "= %lld", result);
	} else {
	    debug(D_DEBUG, "Unknown RPC: %s", line);
            result = -1;
        }

	free(argv);

	snprintf(line, SMQ_LINE_MAX, "%lld\n", result);
	line_size = strlen(line);
	if (link_write(client_link, line, line_size, time(0) + SMQ_Manager_Stalled_Timeout) != line_size) break;
	if (result <= 0) break;
    }
    link_close(client_link);
}

/* Update catalog ----------------------------------------------------------- */

static void
smq_manager_get_information()
{
    struct utsname name;
    char link_path[SMQ_LINE_MAX];

    if (!SMQ_Manager_Owner[0]) 
	if (!username_get(SMQ_Manager_Owner))
	    strncpy(SMQ_Manager_Owner, "unknown", SMQ_LINE_MAX);
    debug(D_DEBUG, "owner = %s", SMQ_Manager_Owner);

    domain_name_cache_guess(SMQ_Manager_Host);
    debug(D_DEBUG, "host = %s", SMQ_Manager_Host);

    if (!SMQ_Manager_Catalog_Port)
	SMQ_Manager_Catalog_Port = datagram_create(0);

    uname(&name);

    string_tolower(name.machine);
    strncpy(SMQ_Manager_CPU, name.machine, SMQ_LINE_MAX);
    debug(D_DEBUG, "cpu = %s", SMQ_Manager_CPU);

    string_tolower(name.sysname);
    strncpy(SMQ_Manager_OpSys, name.sysname, SMQ_LINE_MAX);
    debug(D_DEBUG, "opsys = %s", SMQ_Manager_OpSys);

    string_tolower(name.release);
    strncpy(SMQ_Manager_OpSysVersion, name.release, SMQ_LINE_MAX);
    debug(D_DEBUG, "opsysversion = %s", SMQ_Manager_OpSysVersion);

    /* NOTE: linuxism */
    snprintf(link_path, SMQ_LINE_MAX, "/proc/%d/exe", getpid());
    readlink(link_path, SMQ_Manager_Bin_Path, SMQ_LINE_MAX);
    dirname(SMQ_Manager_Bin_Path);
    debug(D_DEBUG, "bin_path = %s", SMQ_Manager_Bin_Path);
}

static char *
smq_manager_get_message_queue_information(const char *name, int p)
{
    static char text[SMQ_LINE_MAX];
    char  path[SMQ_LINE_MAX];
    int   port;
    int   uptime;
    int	  nmessages;
    struct statfs stats;
    
    port = SMQ_Manager_Port + p;
    snprintf(path, SMQ_LINE_MAX, "%s/%s", SMQ_Manager_Root, name);
    statfs(path, &stats);
    uptime = time(0) - SMQ_Manager_Start_Time;
    nmessages = smq_message_queue_count(SMQ_Manager_Root, name);

    snprintf(text, SMQ_LINE_MAX,
	    "type smq\nversion %s\nsmq_port %d\nowner %s\nsmq_queue %s\nport %d\ncpu %s\nopsys %s\nopsysversion %s\ntotal %lu\navail %lu\nuptime %d\nsmq_messages %d\n",
	    SMQ_VERSION,
	    SMQ_Manager_Port,
	    SMQ_Manager_Owner,
	    name,
	    port,
	    SMQ_Manager_CPU,
	    SMQ_Manager_OpSys,
	    SMQ_Manager_OpSysVersion,
	    stats.f_blocks * stats.f_bsize,
	    stats.f_bavail * stats.f_bsize,
	    uptime,
	    nmessages
    );

    return text;
}

static int
smq_manager_send_catalog_update(const char *queue, int p)
{
    char address[DATAGRAM_ADDRESS_MAX];
    char *text;

    text = smq_manager_get_message_queue_information(queue, p);

    if (domain_name_cache_lookup(CATALOG_HOST, address)) {
	debug(D_DEBUG, "sending %s queue information to %s:%d", queue, CATALOG_HOST, CATALOG_PORT);
	datagram_send(SMQ_Manager_Catalog_Port, text, strlen(text), address, CATALOG_PORT);
    }

    return 1;
}

static void
smq_manager_update_catalog()
{
    char *key;
    void *value;
    int i = 0;
    
    hash_table_firstkey(SMQ_Manager_Exchange_Table);
    while (hash_table_nextkey(SMQ_Manager_Exchange_Table, &key, &value)) 
	smq_manager_send_catalog_update(key, i++);
}

/* Main execution ----------------------------------------------------------- */

int
main(int argc, char *argv[])
{
    struct link *server_link;

    change_process_title_init(argv);
    change_process_title("smq_manager");

    smq_manager_parse_command_line_arguments(argc, argv);

    smq_manager_install_signal_handler(SIGPIPE, smq_manager_ignore_signal);
    smq_manager_install_signal_handler(SIGHUP,  smq_manager_ignore_signal);
    smq_manager_install_signal_handler(SIGXFSZ, smq_manager_ignore_signal);
    smq_manager_install_signal_handler(SIGCHLD, smq_manager_reap_child);
    smq_manager_install_signal_handler(SIGINT,  smq_manager_shutdown_clean);
    smq_manager_install_signal_handler(SIGTERM, smq_manager_shutdown_clean);
    smq_manager_install_signal_handler(SIGQUIT, smq_manager_shutdown_clean);
    smq_manager_install_signal_handler(SIGUSR1, smq_manager_start_exchanges);
    
    smq_manager_get_information(); 
    smq_manager_setup_relay();
    
    SMQ_Manager_Exchange_Table = hash_table_create(0, 0);
    if (!SMQ_Manager_Exchange_Table)
	fatal("could not make exchange table");
    
    server_link = link_serve(SMQ_Manager_Port);
    if (!server_link)
	fatal("could not start server on port %d\n", SMQ_Manager_Port);
    else
	debug(D_DEBUG, "listening on port %d\n", SMQ_Manager_Port);

    SMQ_Manager_Start_Time = time(0);

    while (1) {
	struct link *client_link;
	char address[LINK_ADDRESS_MAX];
	int port;
	pid_t pid;

	if (time(0) > SMQ_Manager_Update_Alarm) {
	    smq_manager_start_exchanges(SIGUSR1);
	    smq_manager_update_catalog();
	    SMQ_Manager_Update_Alarm = time(0) + SMQ_Manager_Update_Timeout;
	}
	
	if (SMQ_Manager_Max_Connections > 0) {
	    if (SMQ_Manager_Total_Connections >= SMQ_Manager_Max_Connections) {
		sleep(5);
		continue;
	    }
	}

	client_link = link_accept(server_link, time(0) + SMQ_Manager_Update_Timeout);
	if (!client_link) continue;

	link_address_remote(client_link, address, &port);
	debug(D_DEBUG, "connected %s:%d", address, port);

	pid = fork();
	if (pid == 0) {
	    change_process_title("smq_manager [connected to %s:%d]", address, port);
	    smq_manager_process_client(client_link);
	    _exit(0);
	} else if (pid > 0) {
	    SMQ_Manager_Total_Connections++;
	    debug(D_DEBUG, "created connection %d (%d total connections)", pid, SMQ_Manager_Total_Connections);
	} else {
	    debug(D_NOTICE, "unable to fork: %s", strerror(errno));
	}
	link_close(client_link);
    }
    
    return EXIT_SUCCESS;
}

/* 
vim: sw=4 sts=4 ts=8 ft=cpp
*/
